import torch
import torch.nn as nn
import torch.nn.functional as F

class upconv(nn.Module):
    def __init__(self, in_dim, out_dim, k_size, scale):
        super(upconv, self).__init__()
        self.scale = scale
        self.conv = conv(in_dim, out_dim, k_size, 1)

    def forward(self, x):
        x = nn.functional.interpolate(x, scale_factor=self.scale, align_corners=True, mode='bilinear')
        return self.conv(x)

class conv(nn.Module):
    def __init__(self, in_dim, out_dim, k_size, stride):
        super(conv, self).__init__()
        self.k_size = k_size
        self.conv = nn.Conv2d(in_dim, out_dim, k_size, stride, padding=(self.k_size - 1) // 2)
        self.bn = nn.BatchNorm2d(out_dim)

    def forward(self, x):
        return F.elu(self.bn(self.conv(x)), inplace=True)


class Desc_head(nn.Module):
    def __init__(self,up_dim=[64, 128, 256, 512]):
        super(Desc_head, self).__init__()

        self.conv_fine = conv(64, 128, 1, 1)
        self.upconv3 = upconv(up_dim[3], up_dim[2], 3, 2)
        self.iconv3 = conv(up_dim[2] + up_dim[2], up_dim[2], 3, 1)

        self.upconv2 = upconv(up_dim[2], up_dim[1], 3, 2)
        self.iconv2 = conv(up_dim[1] + up_dim[1], up_dim[1], 3, 1)

        self.upconv1 = upconv(up_dim[1], up_dim[0], 3, 2)
        self.iconv1 = conv(up_dim[0] + up_dim[0], up_dim[0], 3, 1)

    def skipconnect(self, x1, x2):
        diffY = x2.size()[2] - x1.size()[2]
        diffX = x2.size()[3] - x1.size()[3]
        x1 = F.pad(x1, (diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2))

        x = torch.cat([x2, x1], dim=1)
        return x

    def forward(self,feature):
        x = self.upconv3(feature[3])
        x = self.skipconnect(feature[2], x)
        x = self.iconv3(x)

        x = self.upconv2(x)
        x = self.skipconnect(feature[1], x)
        x = self.iconv2(x)

        x = self.upconv1(x)
        x = self.skipconnect(feature[0], x)
        x = self.iconv1(x)
        x_fine = self.conv_fine(x)  # H/2
        x_fine = F.normalize(x_fine, dim=1)  #做了归一化
        return x_fine
